import random
from PIL import Image
import numpy as np


# 251-255的像素值都会被截断为250
def truncate(img):
    h, w, rgb = img.shape
    ret = np.zeros((h, w, rgb), np.uint8)
    for i in range(h):
        for j in range(w):
            ret[i][j][0] = 250 if img[i][j][0] // 251 else img[i][j][0]
            ret[i][j][1] = 250 if img[i][j][1] // 251 else img[i][j][1]
            ret[i][j][2] = 250 if img[i][j][2] // 251 else img[i][j][2]
    print("done truncating")
    return ret


# 用pixels值与x计算多项式
def cal_polynomial(x, pixels):
    # f_r(x) = pixels[0] + pixels[1]x + ... + pixels[t-1]x^{t-1} mod 251
    ret = pixels[0]
    var = x
    for i in range(1, len(pixels)):
        ret = (ret + pixels[i]*var) % 251
        var = (var * x) % 251
    return ret


# 根据img创建一个宽度为new_w的图片，原本没有图像的部分用0来代替
def create_img(img, new_w):
    h, w, rgb = img.shape
    ret_img = np.zeros((h, new_w, 3), np.uint8)
    for i in range(h):
        for j in range(w):
            ret_img[i][j] = img[i][j]
    print("done creating img")
    return ret_img


# 将图片划分成n个不相互包含的子图（h * (w // k)）
# n: 最多用多少张子图才能合成
# k: 最少用多少张子图才能合成
def split_img(n, k, img):
    h, w, rgb = img.shape
    # shared_w: 子图的宽度
    shared_w = w // k
    if w % k != 0:
        shared_w += 1
        img = create_img(img, shared_w * k)

    ret = []
    # 随机生成n个x值
    x = random.sample(range(1, 251), n)

    for i in range(n):
        print("i:", i, "processing")
        # temp: 子图
        temp = np.zeros((h, shared_w, 3), np.uint8)
        for height in range(h):
            twidth = 0
            for width in range(0, shared_w * k, k):
                rpixels = []
                gpixels = []
                bpixels = []
                for pixel in range(width, width+k):
                    rpixels.append(img[height][pixel][0])
                    gpixels.append(img[height][pixel][1])
                    bpixels.append(img[height][pixel][2])
                R = cal_polynomial(x[i], rpixels)
                G = cal_polynomial(x[i], gpixels)
                B = cal_polynomial(x[i], bpixels)
                temp[height][twidth] = [R, G, B]
                twidth += 1

        im = Image.fromarray(temp)
        # im.show()
        im.save('res/shared_pic_{}.png'.format(i), "png")
        ret.append([temp, x[i]])
    print("done splitting img")
    return ret
